package logger

import (
	"bytes"
	"context"
	"gitee.com/DonHz/basic-go/webook/pkg/logger"
	"github.com/gin-gonic/gin"
	"go.uber.org/atomic"
	"io"
	"time"
)

type MiddlewareLoggerBuilder struct{
	allowReqBody *atomic.Bool
	allowRespBody *atomic.Bool
	logger logger.Logger
	loggerFunc func(ctx context.Context, al *AccessLog)
}

func NewMiddlewareLoggerBuilder(fn func(ctx context.Context, al *AccessLog), l logger.Logger) *MiddlewareLoggerBuilder {
	return &MiddlewareLoggerBuilder{
		allowRespBody: atomic.NewBool(false),
		allowReqBody: atomic.NewBool(false),
		logger: l,
		loggerFunc: fn,
	}
}

func (l *MiddlewareLoggerBuilder) Build() gin.HandlerFunc {
	return func(ctx *gin.Context) {
		start:=time.Now()
		url:= ctx.Request.URL.String()
		if len(url)>1024{
			url=url[:1024]
		}
		al:=&AccessLog{
			Method: ctx.Request.Method,
			URL: url,
		}
		if l.allowReqBody.Load() && ctx.Request.Body!=nil{
			body, err := io.ReadAll(ctx.Request.Body)
			if err != nil {
				l.logger.Error("日志读取请求体失败", logger.Error(err))
				return
			}
			ctx.Request.Body = io.NopCloser(bytes.NewBuffer(body))
			// 这里很消耗CPU和内存，会复制
			al.ReqBody = string(body)
		}
		if l.allowRespBody.Load() {
			ctx.Writer = responseWriter{
				al:al,
				ResponseWriter:ctx.Writer,
			}
		}

		defer func() {
			al.Duration =time.Since(start).String()
			l.loggerFunc(ctx, al)
		}()
		ctx.Next()
	}
}

func (l *MiddlewareLoggerBuilder) AllowReqBody(ok bool) *MiddlewareLoggerBuilder {
	l.allowReqBody.Store(ok)
	return l
}

func (l *MiddlewareLoggerBuilder) AllowRespBody(ok bool) *MiddlewareLoggerBuilder {
	l.allowRespBody.Store(ok)
	return l
}

type responseWriter struct {
	al *AccessLog
	gin.ResponseWriter
}


func (w responseWriter) WriteHeader(statusCode int) {
	w.al.StatusCode=statusCode
	w.ResponseWriter.WriteHeader(statusCode)
}


func (w responseWriter) Write(data []byte) (int, error) {
	w.al.RespBody = string(data)
	return w.ResponseWriter.Write(data)
}

func (w responseWriter) WriteString(data string) (int, error) {
	w.al.RespBody = data
	return w.ResponseWriter.WriteString(data)
}


type AccessLog struct {
	Method string
	URL string
	Duration string
	ReqBody string
	RespBody string
	StatusCode int
}